# RTMP Failover client
#
# Author: Newton Allen
# Date: 12/23/09
#
# Purpose:
#   to facilitate transparent failover of a non-interactive RTMP
#   server serving just videos to clients.
# Note:
#   Server must not send chunks of size > 1 packet (~1452 bytes)
#

import socket
import logging
from cStringIO import StringIO
from collections import namedtuple

import pyamf
import pyamf.amf0
import pyamf.amf3

from rtmp import *

# temp imports
import binascii
import random
from str2int import *
from netutil import stringDump

CONTROL_CSID = 2           # chunk stream ID for chunk stream control packets
COMMAND_CSID = 3           # chunk stream ID for AMF0/3 commands
PLAY_CSID = 8              # chunk stream ID for play() commands
CONNECT_SID = 0            # message stream ID for "connect" command
PLAY_SID = 1               # message stream ID for "play" command
CONNECT_TRANS_ID = 1       # transaction ID for "connect" command
STREAM_TRANS_ID = 2        # transaction ID for "createStream" command
PLAY_TRANS_ID = 0          # transaction ID for "play" command
DEFAULT_CHUNK_SIZE = 128   # default chunk size for new chunk streams
HANDSHAKE_BYTES = 1536     # length of C1, C2 handshake packets

DEFAULT_CONNECT_PARAMS = {'app' : pyamf.Undefined,
                          'flashver' : pyamf.Undefined,
                          'swfUrl' : pyamf.Undefined,
                          'tcUrl' : pyamf.Undefined,
                          'pageUrl' : pyamf.Undefined,
                          'fpad' : False,
                          'audioCodecs' : 3191,
                          'videoCodecs' : 252,
                          'objectEncoding' : pyamf.AMF3,
                          'videoFunction' : 1,
                          'capabilities' : 15}
#DEFAULT_CONNECT_KEYS = set(DEFAULT_CONNECT_PARAMS.keys())
AMF_DECODERS = {MTYPE_AMF0_COMMAND : pyamf.amf0,
                MTYPE_AMF3_COMMAND : pyamf.amf3}

# Example connect packet
"""(Object)#0
  app = "oflaDemo/"
  audioCodecs = 3191
  capabilities = 15
  flashVer = "WIN 10,0,32,18"
  fpad = false
  objectEncoding = 3
  pageUrl = (null)
  swfUrl = (null)
  tcUrl = "rtmp://sns9.cs.princeton.edu/oflaDemo/"
  videoCodecs = 252
  videoFunction = 1"""

# RTMP Chunk received from client
# this is an AMF0/3 Command: receive_state(client_state)
DEFAULT_STATE_PACKET_2 = binascii.unhexlify('030003e500015311000000000002000d726563656976655f737461746500400000000000000005110a0b010f7061676555726c06814766696c653a2f2f2f433a2f55736572732f6e616c6c656e2f4465736b746f702f5363686f6f6c2f49572f466c6578253230576f726b73706163652f54696d656f7574566964656f2f62696e2d64656275672f54696d65c36f7574566964656f2e68746d6c0b746355726c064372746d703a2f2f736e73382e63732e7072696e6365746f6e2e6564752f766f642f076170700607766f6411666c617368766572061d57494e2031302c302c34322c33340b766964656f060d73616d706c650d73776655726c06814566696c653a2f2f2f433a2f5573657273c32f6e616c6c656e2f4465736b746f702f5363686f6f6c2f49572f466c6578253230576f726b73706163652f54696d656f7574566964656f2f62696e2d64656275672f54696d656f7574566964656f2e73776601')

# rtmp://sns9.cs.princeton.edu/oflaDemo/toystory3
DEFAULT_STATE_PACKET = binascii.unhexlify('43000670000095110002000d726563656976655f737461746500000000000000000005110a0b0111666c617368766572061d57494e2031302c302c34322c33340b746355726c064d72746d703a2f2f736e73392e63732e7072696e6365746f6e2e6564752f6f666c6144656d6f2f0b766964656f0613746f7973746f7279330761707006136f666cc36144656d6f2f0974696d6505402b276c8b43958101')
DEFAULT_CONNECT_PACKET = binascii.unhexlify('030000000001021400000000020007636f6e6e656374003ff00000000000000300036170700200096f666c6144656d6f2f0008666c61736856657202000e57494e2031302c302c33322c3138000673776655726c060005746355726c02002672746d703a2f2f736e73392e63732e7072696e6365746f6e2e6564752f6f666c6144656d6f2f00046670616401c300000c6361706162696c697469657300402e000000000000000b617564696f436f646563730040a8ee0000000000000b766964656f436f6465637300406f800000000000000d766964656f46756e6374696f6e003ff000000000000000077061676555726c06000e6f626a656374456e636f64696e6700400800000000000000c30009')

log = logging.getLogger('fo')


'''

1. Ask client for state
2. Get state response and save it
3. Handshake with server
4. connect() to server with params from client
5. createStream() with parameters from client

'''

_netlink_data = 3
def netlink_has_data():
    global _netlink_data
    return _netlink_data > 0

def receive_from_netlink():
    global _netlink_data
    _netlink_data -= 1
    if _netlink_data < 0:
        raise Exception('NETLINK empty')
    return DEFAULT_STATE_PACKET, random.randint(50000, 60000)

def bind_socket_to_port_range(sock, port_start, port_end, max_attempts=10):
    '''
    Bind socket sock to a random port in the range [port_start, port_end)
    Returns the port bound.
    If bind attempt fails max_attempts times, raise a socket.error.
    '''
    for i in range(max_attempts):
        try:
            port = random.randint(port_start, port_end)
            sock.bind(('', port))
            return port
        except socket.error as e:
            err = e
            log.debug("couldn't bind to port %d: %s", port, e)
    raise err

class FailoverError(Exception):
    pass

class ConnectionError(Exception):
    pass

class StateError(Exception):
    pass

class RTMPError(Exception):
    pass

class BufferedSocket:
    "Wrapper around socket to provide read/write semantics"
    
    def __init__(self, sock):
        self.sock = sock
        self.sock.settimeout(12)     # temporary (?) timeout for testing

    def read(self, nbytes, msg=None):
        'Return exactly n bytes from socket.'
        data = ''
        while nbytes > len(data):
            data += self.sock.recv(nbytes - len(data))
        if msg: log.debug(msg)
        return data

    def write(self, data, msg=None):
        'Writes entirety of data to the channel.'
        self.sock.sendall(data)
        if msg: log.debug(msg)


class StreamingFailover:

    def __init__(self, server_addr, streaming_state, port=None, port_range=None, name=None):
        '''
        Create an object to recreate streaming_state on the RTMP server
        at server_addr connecting from local port "port"

        port may be a single port or a tuple representing a range of ports.
        '''
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.buffered_sock = BufferedSocket(self.sock)
        self.rtmp = RTMPStream(self.buffered_sock)
        self.server_addr = server_addr
        self.streaming_state = streaming_state
        self.port = port
        self.port_range = port_range
        self.name = name or ('<StreamingFailover %d>' % id(self))

    def __str__(self):
        return self.name

    def failover(self):
        '''
        Recreate the connection state described by streaming_state on the
        RTMP server at self.server_addr connected from localhost:port
        '''
        try:
            self._failover()
            log.info('%s succeeded!', self)
        except (socket.error, FailoverError) as e:
            log.warning('%s failed: %s', self, e)
        finally:
            pass
            # todo to do: uncomment line below? (or does kernel do this for me?)
            #self.sock.close()

    def _failover(self):
        self._tcpconnect()
        self.rtmp.initiate_handshake()
        self._connect()
        self._create_stream()
        self._play()

        """for i in range(15):
            packet = self.rtmp.read_packet()
            log.debug("received: %s", str(packet))"""

    def _tcpconnect(self):
        '''
        Open TCP connection with RTMP server.

        If self.port is set, connect from that port.  Otherwise, if
        self.port_range is set, connect from a random port in that range.
        '''
        if self.port is not None:
            self.sock.bind(('', self.port))
        elif self.port_range is not None:
            self.port = bind_socket_to_port_range(self.sock, *self.port_range)
        if self.port is not None:
            log.debug('bound to port %d', self.port)
        self.sock.connect(self.server_addr)
        log.debug('connected to server %s', self.server_addr)

    def _connect(self):
        '''
        Send an RTMP-level NetConnection request and wait for _result
        '''
        connect_params = {}
        for k, v in DEFAULT_CONNECT_PARAMS.iteritems():
            connect_params[k] = self.streaming_state.get(k, v)

        connect_header = Header(COMMAND_CSID, 0, MTYPE_AMF0_COMMAND, CONNECT_SID)
        connect_amfobj = ("connect", CONNECT_TRANS_ID, connect_params)
        connect_packet = AMFCommandPacket(connect_header, connect_amfobj)
        self.rtmp.write_packet(connect_packet)

        result = self._wait_for_response(connect_packet).amf_payload
        if result[0] != '_result' or result[3].get('code') != 'NetConnection.Connect.Success':
            raise ConnectionError('%s, %s' % (result[3].get('code'), result[3].get('description')))

        log.debug('connect: success!')

    def _create_stream(self):
        '''
        Send createStream() command to prepare the media stream
        '''
        create_header = Header(COMMAND_CSID, 0, MTYPE_AMF0_COMMAND, CONNECT_SID)
        create_amfobj = ("createStream", STREAM_TRANS_ID, None)
        create_packet = AMFCommandPacket(create_header, create_amfobj)
        self.rtmp.write_packet(create_packet)

        result = self._wait_for_response(create_packet).amf_payload #COMMAND_CSID, STREAM_TRANS_ID)
        if result[0] == "_error":
            raise ConnectionError('createStream: %s' % result[4])

        log.debug('createStream: success!')

    def _play(self):
        '''
        Send play(filename) and seek(time) commands to begin the streaming
        '''
        play_header = Header(PLAY_CSID, 0, MTYPE_AMF3_COMMAND, PLAY_SID)
        play_params = ("play", PLAY_TRANS_ID, None, self.streaming_state.get('video'))
        play_packet = AMFCommandPacket(play_header, play_params)
        seek_params = ("seek", 0, None, self.streaming_state.get('time') * 1000)
        seek_packet = AMFCommandPacket(play_header, seek_params)

        self.rtmp.write_packet(play_packet, seek_packet)

        log.debug('play: sent play(%s); seek(%d msec.)',
                  self.streaming_state.get('video'),
                  self.streaming_state.get('time') * 1000)

    def _wait_for_response(self, amf_command_packet):
        '''
        Reads and discards packets until the AMF command response packet
        to amf_command_packet.  Returns the response packet
        '''

        csid = amf_command_packet.header.csid
        trans_id = amf_command_packet.amf_payload[1]
        log.debug('wait_for_response: waiting on ID #%d', trans_id)

        while 1:
            try:
                packet = self.rtmp.read_packet()
                """log.debug('wait_for_response: received packet:\n%s\n%s',
                          binascii.hexlify(packet.data), stringDump(packet.data))
                log.debug('packet type = %s', packet.__class__)
                log.debug('header = %s', packet.header)
                if packet.__class__ is AMFCommandPacket:
                    log.debug('payload = %s', packet.amf_payload)"""
                if packet.__class__ is AMFCommandPacket and packet.header.csid == csid and\
                   len(packet.amf_payload) >= 4 and packet.amf_payload[1] == trans_id:
                    """if len(packet.amf_payload) < 4:
                        raise ConnectionError('connect: malformed _result packet')"""
                    break
            except pyamf.DecodeError as e:
                log.warning('wait_for_response: AMF decode error')

        return packet

class ChunkStream:
    '''
    Holds the state for a single chunk stream of RTMP.
    All messages in this stream have the same chunk stream ID ("csid")
    '''

    def __init__(self, csid):
        self.csid = csid
        self.time = 0
        self.dtime = 0
        self.msg_len = 0
        self.remaining_msg_len = 0
        self.mtype = 0
        self.sid = 0

    def __str__(self):
        return 'ChunkStream(csid=%s, time=%s, dtime=%s, msg_len=%s, mtype=%s, sid=%s)' %\
               (self.csid, self.time, self.dtime, self.msg_len, self.mtype, self.sid)

# Shortcut class definition
Header = namedtuple('Header', 'csid time mtype sid')

class Packet:
    def __init__(self, header, data):
        self.header = header
        self.data = data

    def __str__(self):
        return '%s :\n%s' % (str(self.header), stringDump(self.data))

    @staticmethod
    def make_packet(header, data):
        '''
        Initializes and returns a packet as above, *but* if a more
        specialized packet type exists (e.g. AMF0CommandPacket),
        returns one of those instead.
        '''
        # If AMF packet, decode contents
        if header.mtype in (MTYPE_AMF0_COMMAND, MTYPE_AMF3_COMMAND):
            return AMFCommandPacket(header, data=data)

        return Packet(header, data)

class AMFCommandPacket:
    def __init__(self, header, amf_payload=None, data=None):
        assert header.mtype in (MTYPE_AMF0_COMMAND, MTYPE_AMF3_COMMAND)
        self.header = header
        self.__amf_payload = amf_payload
        self.__data = data

    def __str__(self):
        return '%s :\n%s' % (str(self.header), self.amf_payload)

    def get_data(self):
        if not self.__data:
            self.__data = pyamf.amf0.encode(*self.__amf_payload).getvalue()
            if self.header.mtype == MTYPE_AMF3_COMMAND:
                # prepending '\x00' is undocumented, but necessary
                self.__data = '\x00' + self.__data
        return self.__data

    def get_amf_payload(self):
        if not self.__amf_payload:
            data = self.__data
            # remove '\x00' here? i.e. self.__data[1:]
            if self.header.mtype == MTYPE_AMF3_COMMAND:
                data = data[1:]
            self.__amf_payload = tuple(pyamf.amf0.decode(data))
        return self.__amf_payload

    data = property(get_data)
    amf_payload = property(get_amf_payload)

class RTMPStream:
    '''
    A two-way RTMP connection to parse and build RTMP packets
    Owns an iobuf, which must support read(nbytes) and write(string) operations.
    '''

    def __init__(self, iobuf):
        self.iobuf = iobuf
        self.chunk_size_in = DEFAULT_CHUNK_SIZE   # chunk size for received packets
        self.chunk_size_out = DEFAULT_CHUNK_SIZE  # chunk size for sent packets
        self.chunk_streams = {}

    def initiate_handshake(self):
        '''
        Connect to an RTMP server over iobuf
        '''
        read = self.iobuf.read
        write = self.iobuf.write
        
        c0 = RTMP_VERSION
        c1 = '\x00' * 8 + '\x55' * (HANDSHAKE_BYTES - 8)
        write(c0 + c1, 'handshake: sent c0 and c1')

        s0 = read(1)
        s1 = read(HANDSHAKE_BYTES, 'handshake: received s0 and s1')
        if s0 != RTMP_VERSION:
            raise RTMPError('handshake: invalid RTMP version')

        c2 = s1
        write(c2, 'handshake: sent c2')

        s2 = read(HANDSHAKE_BYTES, 'handshake: received s2')
        if s2[8:] != c1[8:]:
            raise RTMPError('handshake: non-matching reciprocal handshake')

        log.debug('handshake: complete!')

    def get_stream(self, csid):
        if csid not in self.chunk_streams:
            self.chunk_streams[csid] = ChunkStream(csid)
        return self.chunk_streams[csid]

    def parse_packet(self, read):
        '''
        Returns a packet object representing the data returned by read(nbytes)
        '''
        basic_hdr = ord(read(1))
        fmt = basic_hdr >> 6                # format of chunk stream header
        csid = basic_hdr & 0x3F             # chunk stream ID
        cs = self.get_stream(csid)

        if 0 < fmt < 3:
            cs.dtime = uint24BE(read(3))    # relative timestamp
        if fmt == 0:
            cs.time = uint24BE(read(3))     # absolute timestamp
        else:
            cs.time += cs.dtime
        if fmt < 2:
            cs.msg_len = uint24BE(read(3))  # message length
            cs.mtype = ord(read(1))         # message type ID
        if fmt == 0:
            cs.sid = uint32LE(read(4))      # message stream ID

        # unchunkify payload
        internal_hdr_len = cs.msg_len / self.chunk_size_in
        chunked_payload = read(cs.msg_len + internal_hdr_len)
        payload = self.unchunk(chunked_payload)

        if cs.mtype == MTYPE_SET_CHUNK_SIZE:
            self.chunk_size_in = uint32BE(payload)
            log.debug('parse_packet: set chunk size to %d', self.chunk_size_in)

        header = Header(csid, cs.time, cs.mtype, cs.sid)
        packet = Packet.make_packet(header, payload)

        return packet

    def parse_chunk(self, read):
        '''
        Returns the next chunk from the channel abstracted by read(nbytes).
        This may or may not be a full packet.

        This is actually the correct command to use -- parse_packet() usually
        works, but not if chunks from multiple chunk streams and with compressed
        headers are interleaved.
        '''
        basic_hdr = ord(read(1))
        fmt = basic_hdr >> 6                # format of chunk stream header
        csid = basic_hdr & 0x3F             # chunk stream ID
        cs = self.get_stream(csid)

        if cs.remaining_msg_len == 0:
            if 0 < fmt < 3:
                cs.dtime = uint24BE(read(3))    # relative timestamp
            if fmt == 0:
                cs.time = uint24BE(read(3))     # absolute timestamp
            else:
                cs.time += cs.dtime
            if fmt < 2:
                cs.msg_len = uint24BE(read(3))  # message length
                cs.mtype = ord(read(1))         # message type ID
            if fmt == 0:
                cs.sid = uint32LE(read(4))      # message stream ID
            cs.remaining_msg_len = cs.msg_len
        else:
            assert fmt == 3
            cs.time += cs.dtime

        incremental_msg_len = min(cs.remaining_msg_len, self.chunk_size_in)
        cs.remaining_msg_len -= incremental_msg_len

        chunk = read(incremental_msg_len)

        if cs.mtype == MTYPE_SET_CHUNK_SIZE:
            self.chunk_size_in = uint32BE(chunk)
            log.debug('parse_packet: set chunk size to %d', self.chunk_size_in)

        header = Header(csid, cs.time, cs.mtype, cs.sid)
        packet = Packet.make_packet(header, chunk)

        return packet

    def build_packet(self, packet):
        '''
        Return a string encoding the desired RTMP packet.
        Header compression is not used.
        '''
        header = packet.header
        msg_len = len(packet.data)

        string = chr(header.csid) + str24BE(header.time) + str24BE(msg_len) +\
                 chr(header.mtype) + str32LE(header.sid) + self.chunk(packet.data, header.csid)

        return string

    def read_packet(self):
        '''Read the next packet from the RTMP stream'''
        return self.parse_packet(self.iobuf.read)

    def read_chunk(self):
        '''Read the next chunk from the RTMP stream'''
        return self.parse_chunk(self.iobuf.read)    

    def write_packet(self, *packets):
        '''Build a packet (or several) and write it to iobuf'''
        output = ''.join([self.build_packet(packet) for packet in packets])
        self.iobuf.write(output)

    def unchunk(self, data):
        # Remove 1-byte header that occurs after each chunk of bytes
        return ''.join([data[i:i + self.chunk_size_in] for i in range(0, len(data), self.chunk_size_in + 1)])

    def chunk(self, data, csid):
        # Insert 1-byte header after each chunk of bytes
        return chr(0xC0 + csid).join([data[i:i + self.chunk_size_out] for i in range(0, len(data), self.chunk_size_out)])

def decode_state(streaming_state_string):
    '''
    Convert a raw (binary data) streaming state string into a
    dictionary containing the streaming state information
    '''
    rtmp = RTMPStream(StringIO(streaming_state_string))
    state_packet = rtmp.read_packet()
    if state_packet.__class__ != AMFCommandPacket:
        return None
    try:
        state = state_packet.amf_payload[3]
    except (pyamf.DecodeError, IndexError) as e:
        return None   # not a streaming state packet

    return state

if __name__ == '__main__':
    failover_count = 0

    logging.basicConfig(level=logging.DEBUG,
                        format='%(asctime)s.%(msecs)03d %(lineno)d %(levelname)s %(message)s',
                        datefmt='%H:%M:%S')
    
    while 1:
        failover_count += 1
        failover_name = 'Failover #%d' % failover_count
        
        get_state_response, port = receive_from_netlink()
        log.info('starting %s on port %d', failover_name, port)

        # get state from state_packet
        try:
            streaming_state = decode_state(get_state_response)
        except (StateError, pyamf.DecodeError) as e:
            log.warning('aborting failover: streaming state %s' % e)
            continue

        # do failover
        f = StreamingFailover(('localhost', RTMP_PORT), streaming_state, port, failover_name)
        f.failover()

        log.info('goodbye')
        break

